# 导入通义大模型
from langchain_community.llms import Tongyi
# 导入模板类
from langchain.prompts import ChatPromptTemplate
from langchain.schema import SystemMessage, HumanMessage, AIMessage

# sy = SystemMessage(content="你是一个起名大师. 你的名字叫宋大师.")
# hu = HumanMessage(content="我家是宝，请起3个好听的名字？")
# message = [sy, hu]


def generate_name(species, trait, nameStyle):
        # 定义模板
        message = [
                ("system", "假设你是一位养宠物的大师，非常擅长给宠物起名字"),
                ("human", "我家宠物是一只{species}，有一些特征，如:{trait}，名字有{nameStyle}风格，请帮我起3个好听的名字吧"),
        ]

        # 实例化模板类
        chartmp = ChatPromptTemplate.from_messages(message)
        prompt = chartmp.format_messages(species=species, trait=trait, nameStyle=nameStyle)
        print(prompt)

        # 实例化通义大模型
        tongyi = Tongyi()
        ret = tongyi.invoke(prompt)
        print(ret)
        # print(type(ret))
        return ret

if __name__ == '__main__':
    generate_name('猫', '浑身白色，头顶有些黑色', '霸气一点')